# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: donggu

import pandas as pd
import msgpack
import matplotlib.pyplot as plt

# matplotlib.style.use('default')

plt.rcParams['figure.figsize'] = 12, 8
plt.rcParams['font.size'] = 9
plt.rcParams['legend.fontsize'] = 9
plt.rcParams['xtick.labelsize'] = 9
plt.rcParams['ytick.labelsize'] = 9

path = 'feed-dump.msgpack'
with open(path, 'rb') as f:
  obj = msgpack.load(f)

for (product, data) in obj[b'products'].items():
  product = str(product)

  trade_timestamps = pd.to_datetime(data[b'trade_timestamps'], unit='ns', utc=True)
  book_timestamps = pd.to_datetime(data[b'book_timestamps'], unit='ns', utc=True)

  assert (len(trade_timestamps) == len(data[b'trade_prices']))
  assert (len(book_timestamps) == len(data[b'ask0']))
  assert (len(book_timestamps) == len(data[b'bid0']))

  plt.scatter(trade_timestamps,
              data[b'trade_prices'],
              s=0.2,
              label=product + ':trade (%d)' % len(trade_timestamps))
  # plt.plot(book_timestamps, data[b'ask0'], drawstyle='steps-post', label=product + ':ask0')
  # plt.plot(book_timestamps, data[b'bid0'], drawstyle='steps-post', label=product + ':bid0')
  plt.fill_between(book_timestamps,
                   data[b'bid0'],
                   data[b'ask0'],
                   alpha=0.5,
                   lw=0.5,
                   label=product + ':ask/bid (%d)' % len(book_timestamps))

# plt.show()
plt.minorticks_on()
plt.xlabel('Price')
plt.ylabel('Time')
plt.legend()

#plt.autoscale()  # not working
start, end = pd.to_datetime(obj[b'interval'], unit='ns', utc=True)
plt.xlim(start, end)

plt.title('Feed ' + str(start) + ' ~ ' + str(end))

plt.savefig('feed-dump.png')
